#!/usr/bin/env python3
# 版本: 0.5.0.20240803
# 作者: 陈振玺
# 功能: 服务器任务调度、资源监控和Webhook推送工具

import os
import sys
import time
import json
import psutil
import subprocess
import threading
import requests
import argparse
import tempfile
from datetime import datetime
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass, asdict
from concurrent.futures import ThreadPoolExecutor


@dataclass
class ResourceUsage:
    """资源使用情况数据类"""
    timestamp: str
    cpu_percent: float
    memory_mb: float
    memory_percent: float
    disk_io_read_mb: float
    disk_io_write_mb: float
    network_sent_mb: float
    network_recv_mb: float
    process_count: int
    load_average: Tuple[float, float, float]


@dataclass
class TaskStatus:
    """任务状态数据类"""
    task_id: str
    command: str
    status: str  # running, completed, failed
    start_time: str
    end_time: Optional[str]
    exit_code: Optional[int]
    pid: Optional[int]
    resource_usage: List[ResourceUsage]


class ResourceMonitor:
    """资源监控器"""
    
    def __init__(self):
        self.monitoring = False
        self.process_pids = set()
        self.initial_disk_io = None
        self.initial_network_io = None
    
    def add_process(self, pid: int):
        """添加需要监控的进程PID"""
        self.process_pids.add(pid)
    
    def remove_process(self, pid: int):
        """移除监控的进程PID"""
        self.process_pids.discard(pid)
    
    def get_process_tree_pids(self, parent_pid: int) -> List[int]:
        """获取进程树中所有子进程的PID"""
        try:
            parent = psutil.Process(parent_pid)
            children = parent.children(recursive=True)
            return [parent_pid] + [child.pid for child in children]
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            return []
    
    def get_process_resource_usage(self, pids: List[int]) -> Tuple[float, float]:
        """获取指定进程列表的CPU和内存使用情况"""
        total_cpu = 0.0
        total_memory = 0.0
        
        for pid in pids:
            try:
                process = psutil.Process(pid)
                total_cpu += process.cpu_percent()
                total_memory += process.memory_info().rss / 1024 / 1024  # MB
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        
        return total_cpu, total_memory
    
    def get_system_resource_usage(self) -> ResourceUsage:
        """获取系统资源使用情况"""
        # CPU使用率
        cpu_percent = psutil.cpu_percent(interval=1)
        
        # 内存使用情况
        memory = psutil.virtual_memory()
        memory_mb = memory.used / 1024 / 1024
        memory_percent = memory.percent
        
        # 磁盘IO
        disk_io = psutil.disk_io_counters()
        if self.initial_disk_io is None:
            self.initial_disk_io = disk_io
            disk_read_mb = 0
            disk_write_mb = 0
        else:
            disk_read_mb = (disk_io.read_bytes - self.initial_disk_io.read_bytes) / 1024 / 1024
            disk_write_mb = (disk_io.write_bytes - self.initial_disk_io.write_bytes) / 1024 / 1024
        
        # 网络IO
        network_io = psutil.net_io_counters()
        if self.initial_network_io is None:
            self.initial_network_io = network_io
            network_sent_mb = 0
            network_recv_mb = 0
        else:
            network_sent_mb = (network_io.bytes_sent - self.initial_network_io.bytes_sent) / 1024 / 1024
            network_recv_mb = (network_io.bytes_recv - self.initial_network_io.bytes_recv) / 1024 / 1024
        
        # 进程数量
        process_count = len(psutil.pids())
        
        # 系统负载
        load_average = os.getloadavg()
        
        return ResourceUsage(
            timestamp=datetime.now().isoformat(),
            cpu_percent=cpu_percent,
            memory_mb=memory_mb,
            memory_percent=memory_percent,
            disk_io_read_mb=disk_read_mb,
            disk_io_write_mb=disk_write_mb,
            network_sent_mb=network_sent_mb,
            network_recv_mb=network_recv_mb,
            process_count=process_count,
            load_average=load_average
        )


class WebhookNotifier:
    """Webhook通知器"""
    
    def __init__(self, webhook_url: str):
        self.webhook_url = webhook_url
    
    def send_notification(self, message: str, msg_type: str = "text") -> bool:
        """发送通知到webhook"""
        try:
            payload = {
                "msg_type": msg_type,
                "content": {
                    "text": message
                }
            }
            
            response = requests.post(
                self.webhook_url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=10
            )
            
            if response.status_code == 200:
                print(f"✅ Webhook通知发送成功: {message[:50]}...")
                return True
            else:
                print(f"❌ Webhook通知发送失败: {response.status_code} - {response.text}")
                return False
                
        except Exception as e:
            print(f"❌ Webhook通知发送异常: {str(e)}")
            return False
    
    def send_task_status(self, task: TaskStatus) -> bool:
        """发送任务状态通知"""
        if task.status == "running":
            message = f"🚀 任务开始执行\n" \
                     f"任务ID: {task.task_id}\n" \
                     f"命令: {task.command}\n" \
                     f"开始时间: {task.start_time}\n" \
                     f"进程PID: {task.pid}"
        elif task.status == "completed":
            duration = self._calculate_duration(task.start_time, task.end_time)
            message = f"✅ 任务执行完成\n" \
                     f"任务ID: {task.task_id}\n" \
                     f"命令: {task.command}\n" \
                     f"执行时长: {duration}\n" \
                     f"退出码: {task.exit_code}"
        elif task.status == "failed":
            duration = self._calculate_duration(task.start_time, task.end_time)
            message = f"❌ 任务执行失败\n" \
                     f"任务ID: {task.task_id}\n" \
                     f"命令: {task.command}\n" \
                     f"执行时长: {duration}\n" \
                     f"退出码: {task.exit_code}"
        else:
            message = f"📊 任务状态更新\n" \
                     f"任务ID: {task.task_id}\n" \
                     f"状态: {task.status}"
        
        return self.send_notification(message)
    
    def send_resource_report(self, tasks: List[TaskStatus], current_usage: ResourceUsage) -> bool:
        """发送资源使用报告"""
        running_tasks = [t for t in tasks if t.status == "running"]
        completed_tasks = [t for t in tasks if t.status == "completed"]
        failed_tasks = [t for t in tasks if t.status == "failed"]
        
        message = f"📊 系统资源监控报告\n" \
                 f"时间: {current_usage.timestamp}\n" \
                 f"\n🖥️ 系统资源:\n" \
                 f"CPU使用率: {current_usage.cpu_percent:.1f}%\n" \
                 f"内存使用: {current_usage.memory_mb:.1f}MB ({current_usage.memory_percent:.1f}%)\n" \
                 f"磁盘IO: 读取{current_usage.disk_io_read_mb:.1f}MB, 写入{current_usage.disk_io_write_mb:.1f}MB\n" \
                 f"网络IO: 发送{current_usage.network_sent_mb:.1f}MB, 接收{current_usage.network_recv_mb:.1f}MB\n" \
                 f"进程数: {current_usage.process_count}\n" \
                 f"系统负载: {current_usage.load_average[0]:.2f}, {current_usage.load_average[1]:.2f}, {current_usage.load_average[2]:.2f}\n" \
                 f"\n📋 任务状态:\n" \
                 f"运行中: {len(running_tasks)}个\n" \
                 f"已完成: {len(completed_tasks)}个\n" \
                 f"失败: {len(failed_tasks)}个"
        
        return self.send_notification(message)
    
    def _calculate_duration(self, start_time: str, end_time: Optional[str]) -> str:
        """计算任务执行时长"""
        if not end_time:
            return "未知"
        
        try:
            start = datetime.fromisoformat(start_time)
            end = datetime.fromisoformat(end_time)
            duration = end - start
            
            total_seconds = int(duration.total_seconds())
            hours = total_seconds // 3600
            minutes = (total_seconds % 3600) // 60
            seconds = total_seconds % 60
            
            if hours > 0:
                return f"{hours}小时{minutes}分钟{seconds}秒"
            elif minutes > 0:
                return f"{minutes}分钟{seconds}秒"
            else:
                return f"{seconds}秒"
        except Exception:
            return "未知"


class TaskScheduler:
    """任务调度器"""
    
    def __init__(self, webhook_url: str, monitor_interval: int = 30):
        self.webhook_url = webhook_url
        self.monitor_interval = monitor_interval
        self.tasks: List[TaskStatus] = []
        self.resource_monitor = ResourceMonitor()
        self.webhook_notifier = WebhookNotifier(webhook_url)
        self.running = False
        self.monitor_thread = None
    
    def load_commands(self, script_file: str) -> List[str]:
        """从脚本文件加载命令列表"""
        try:
            with open(script_file, 'r', encoding='utf-8') as f:
                commands = [line.strip() for line in f.readlines() if line.strip() and not line.startswith('#')]
            print(f"📋 成功加载 {len(commands)} 个命令")
            return commands
        except Exception as e:
            print(f"❌ 加载脚本文件失败: {str(e)}")
            return []
    
    def execute_command(self, command: str, task_id: str) -> TaskStatus:
        """执行单个命令"""
        task = TaskStatus(
            task_id=task_id,
            command=command,
            status="running",
            start_time=datetime.now().isoformat(),
            end_time=None,
            exit_code=None,
            pid=None,
            resource_usage=[]
        )
        
        temp_script_path = None
        
        try:
            print(f"🚀 开始执行任务 {task_id}: {command}")
            
            # 创建临时脚本文件
            temp_script_fd, temp_script_path = tempfile.mkstemp(suffix='.sh', prefix=f'{task_id}_')
            
            # 写入命令到临时脚本文件
            with os.fdopen(temp_script_fd, 'w') as temp_file:
                temp_file.write("#!/bin/bash\n")
                temp_file.write(f"# 任务ID: {task_id}\n")
                temp_file.write(f"# 生成时间: {datetime.now().isoformat()}\n")
                temp_file.write(f"# 原始命令: {command}\n")
                temp_file.write(f"{command}\n")
            
            # 设置脚本文件可执行权限
            os.chmod(temp_script_path, 0o755)
            
            print(f"📝 命令已写入临时脚本: {temp_script_path}")
            
            # 启动进程执行脚本文件
            process = subprocess.Popen(
                ['/bin/bash', temp_script_path],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            
            task.pid = process.pid
            self.resource_monitor.add_process(process.pid)
            
            # 发送任务开始通知
            self.webhook_notifier.send_task_status(task)
            
            # 等待进程完成
            stdout, stderr = process.communicate()
            
            task.end_time = datetime.now().isoformat()
            task.exit_code = process.returncode
            
            if process.returncode == 0:
                task.status = "completed"
                print(f"✅ 任务 {task_id} 执行成功")
            else:
                task.status = "failed"
                print(f"❌ 任务 {task_id} 执行失败，退出码: {process.returncode}")
                if stderr:
                    print(f"错误信息: {stderr}")
            
            # 发送任务完成通知
            self.webhook_notifier.send_task_status(task)
            
        except Exception as e:
            task.status = "failed"
            task.end_time = datetime.now().isoformat()
            print(f"❌ 任务 {task_id} 执行异常: {str(e)}")
            self.webhook_notifier.send_task_status(task)
        
        finally:
            if task.pid:
                self.resource_monitor.remove_process(task.pid)
            
            # 清理临时脚本文件
            if temp_script_path and os.path.exists(temp_script_path):
                try:
                    os.unlink(temp_script_path)
                    print(f"🗑️ 临时脚本已清理: {temp_script_path}")
                except Exception as e:
                    print(f"⚠️ 清理临时脚本失败: {str(e)}")
        
        return task
    
    def start_monitoring(self):
        """启动资源监控"""
        self.running = True
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.monitor_thread.start()
        print(f"📊 资源监控已启动，间隔: {self.monitor_interval}秒")
    
    def stop_monitoring(self):
        """停止资源监控"""
        self.running = False
        if self.monitor_thread:
            self.monitor_thread.join()
        print("📊 资源监控已停止")
    
    def _monitor_loop(self):
        """监控循环"""
        while self.running:
            try:
                # 获取当前资源使用情况
                current_usage = self.resource_monitor.get_system_resource_usage()
                
                # 发送资源报告
                self.webhook_notifier.send_resource_report(self.tasks, current_usage)
                
                # 等待下次监控
                time.sleep(self.monitor_interval)
                
            except Exception as e:
                print(f"❌ 监控循环异常: {str(e)}")
                time.sleep(5)  # 异常时短暂等待
    
    def run_tasks(self, script_file: str, max_workers: int = 1):
        """运行任务脚本"""
        commands = self.load_commands(script_file)
        if not commands:
            print("❌ 没有可执行的命令")
            return
        
        # 启动监控
        self.start_monitoring()
        
        # 发送开始通知
        start_message = f"🎯 任务调度开始\n" \
                       f"脚本文件: {script_file}\n" \
                       f"命令数量: {len(commands)}\n" \
                       f"并发数: {max_workers}\n" \
                       f"监控间隔: {self.monitor_interval}秒"
        self.webhook_notifier.send_notification(start_message)
        
        try:
            # 使用线程池执行任务
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                futures = []
                
                for i, command in enumerate(commands, 1):
                    task_id = f"task_{i:03d}"
                    future = executor.submit(self.execute_command, command, task_id)
                    futures.append(future)
                
                # 等待所有任务完成
                for future in futures:
                    task = future.result()
                    self.tasks.append(task)
        
        except KeyboardInterrupt:
            print("\n⚠️ 收到中断信号，正在停止任务...")
        
        finally:
            # 停止监控
            self.stop_monitoring()
            
            # 发送最终报告
            self._send_final_report()
    
    def _send_final_report(self):
        """发送最终执行报告"""
        completed_tasks = [t for t in self.tasks if t.status == "completed"]
        failed_tasks = [t for t in self.tasks if t.status == "failed"]
        
        total_duration = self._calculate_total_duration()
        
        message = f"🏁 任务调度完成\n" \
                 f"总任务数: {len(self.tasks)}\n" \
                 f"成功: {len(completed_tasks)}个\n" \
                 f"失败: {len(failed_tasks)}个\n" \
                 f"总耗时: {total_duration}\n" \
                 f"完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
        
        self.webhook_notifier.send_notification(message)
    
    def _calculate_total_duration(self) -> str:
        """计算总执行时长"""
        if not self.tasks:
            return "0秒"
        
        start_times = [datetime.fromisoformat(t.start_time) for t in self.tasks if t.start_time]
        end_times = [datetime.fromisoformat(t.end_time) for t in self.tasks if t.end_time]
        
        if not start_times or not end_times:
            return "未知"
        
        total_start = min(start_times)
        total_end = max(end_times)
        duration = total_end - total_start
        
        total_seconds = int(duration.total_seconds())
        hours = total_seconds // 3600
        minutes = (total_seconds % 3600) // 60
        seconds = total_seconds % 60
        
        if hours > 0:
            return f"{hours}小时{minutes}分钟{seconds}秒"
        elif minutes > 0:
            return f"{minutes}分钟{seconds}秒"
        else:
            return f"{seconds}秒"


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="服务器任务调度、资源监控和webhook推送工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  task-monitor -s script.sh
  task-monitor -s script.sh -w https://your-webhook-url
  task-monitor -s script.sh -j 4 -i 60
        """
    )
    
    parser.add_argument(
        "-s", "--script",
        required=True,
        help="包含bash命令的脚本文件路径"
    )
    
    parser.add_argument(
        "-w", "--webhook",
        default="https://open.feishu.cn/open-apis/bot/v2/hook/4ed84c9a-e927-4c6d-984f-7990cab894bd",
        help="Webhook URL（默认为飞书机器人）"
    )
    
    parser.add_argument(
        "-j", "--jobs",
        type=int,
        default=1,
        help="并发执行的任务数量（默认: 1）"
    )
    
    parser.add_argument(
        "-i", "--interval",
        type=int,
        default=30,
        help="资源监控推送间隔（秒，默认: 30）"
    )
    
    args = parser.parse_args()
    
    # 检查脚本文件是否存在
    if not os.path.isfile(args.script):
        print(f"❌ 脚本文件不存在: {args.script}")
        sys.exit(1)
    
    print(f"🎯 任务监控工具启动")
    print(f"📄 脚本文件: {args.script}")
    print(f"🔗 Webhook URL: {args.webhook}")
    print(f"⚡ 并发数: {args.jobs}")
    print(f"⏰ 监控间隔: {args.interval}秒")
    print("-" * 50)
    
    # 创建任务调度器并运行
    scheduler = TaskScheduler(args.webhook, args.interval)
    scheduler.run_tasks(args.script, args.jobs)
    
    print("\n🏁 任务调度完成")


if __name__ == "__main__":
    main()